import os
import re
import zipfile
import uuid
from lxml import etree
from PIL import Image
from io import BytesIO
import argparse
import shutil


def extract_all_docx_images(docx_path, output_folder):
    os.makedirs(output_folder, exist_ok=True)

    temp_dir = os.path.join(output_folder, "temp_extract")
    os.makedirs(temp_dir, exist_ok=True)

    with zipfile.ZipFile(docx_path, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)

    image_paths = []

    document_xml_path = os.path.join(temp_dir, "word", "document.xml")
    image_paths += extract_images_from_xml(document_xml_path, temp_dir, output_folder, "body")

    headers_dir = os.path.join(temp_dir, "word", "header")
    if os.path.exists(headers_dir):
        for header_file in os.listdir(headers_dir):
            if header_file.endswith(".xml"):
                header_path = os.path.join(headers_dir, header_file)
                image_paths += extract_images_from_xml(header_path, temp_dir, output_folder, "header")

    footers_dir = os.path.join(temp_dir, "word", "footer")
    if os.path.exists(footers_dir):
        for footer_file in os.listdir(footers_dir):
            if footer_file.endswith(".xml"):
                footer_path = os.path.join(footers_dir, footer_file)
                image_paths += extract_images_from_xml(footer_path, temp_dir, output_folder, "footer")

    rels_path = os.path.join(temp_dir, "word", "_rels", "document.xml.rels")
    image_paths += extract_embedded_images(rels_path, temp_dir, output_folder)

    shutil.rmtree(temp_dir)

    print(f"共提取 {len(image_paths)} 张图片")
    return image_paths


def extract_images_from_xml(xml_path, temp_dir, output_folder, context):
    if not os.path.exists(xml_path):
        return []

    image_paths = []

    parser = etree.XMLParser(recover=True)
    tree = etree.parse(xml_path, parser)
    root = tree.getroot()

    namespaces = {
        'a': 'http://schemas.openxmlformats.org/drawingml/2006/main',
        'wp': 'http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing',
        'r': 'http://schemas.openxmlformats.org/officeDocument/2006/relationships'
    }

    for drawing in root.xpath('//wp:docPr', namespaces=namespaces):
        blip = drawing.xpath('.//a:blip', namespaces=namespaces)
        if blip:
            r_id = blip[0].get('{{{}}}embed'.format(namespaces['r']))
            if r_id:
                image_path = get_image_path_from_rels(r_id, xml_path, temp_dir)
                if image_path:
                    saved_path = save_image(image_path, output_folder, context)
                    if saved_path:
                        image_paths.append(saved_path)

    return image_paths


def extract_embedded_images(rels_path, temp_dir, output_folder):
    if not os.path.exists(rels_path):
        return []

    image_paths = []

    parser = etree.XMLParser(recover=True)
    tree = etree.parse(rels_path, parser)
    root = tree.getroot()

    namespaces = {
        'r': 'http://schemas.openxmlformats.org/package/2006/relationships'
    }

    for rel in root.xpath('//r:Relationship', namespaces=namespaces):
        if "image" in rel.get('Type'):
            r_id = rel.get('Id')
            target = rel.get('Target')
            if target:
                image_path = os.path.join(temp_dir, "word", target.replace('../', ''))
                if os.path.exists(image_path):
                    saved_path = save_image(image_path, output_folder, "embedded")
                    if saved_path:
                        image_paths.append(saved_path)

    return image_paths


def get_image_path_from_rels(r_id, xml_path, temp_dir):
    rels_dir = os.path.join(os.path.dirname(xml_path), "_rels")
    rels_file = os.path.basename(xml_path) + ".rels"
    rels_path = os.path.join(rels_dir, rels_file)

    if not os.path.exists(rels_path):
        return None

    parser = etree.XMLParser(recover=True)
    tree = etree.parse(rels_path, parser)
    root = tree.getroot()

    namespaces = {
        'r': 'http://schemas.openxmlformats.org/package/2006/relationships'
    }

    for rel in root.xpath('//r:Relationship', namespaces=namespaces):
        if rel.get('Id') == r_id:
            target = rel.get('Target')
            if target:
                image_path = os.path.join(temp_dir, "word", target.replace('../', ''))
                if os.path.exists(image_path):
                    return image_path
    return None


def save_image(image_path, output_folder, context):
    with open(image_path, 'rb') as f:
        img_data = f.read()

    try:
        img = Image.open(BytesIO(img_data))
        img_format = img.format
        img.close()
    except:
        ext = os.path.splitext(image_path)[1].lower()
        if ext == ".png":
            img_format = "PNG"
        elif ext in [".jpg", ".jpeg"]:
            img_format = "JPEG"
        elif ext == ".gif":
            img_format = "GIF"
        elif ext == ".bmp":
            img_format = "BMP"
        elif ext == ".tiff":
            img_format = "TIFF"
        else:
            img_format = "PNG"

    file_ext = f".{img_format.lower()}"
    img_name = f"{context}_{uuid.uuid4().hex}{file_ext}"
    output_path = os.path.join(output_folder, img_name)

    with open(output_path, 'wb') as f:
        f.write(img_data)

    return output_path


def main():


    path="D:\\hbyt\\AI智能投标\\能力组件-附件\\ITO行业发展趋势_v1.0_2211.docx"
    out_path="D:\\hbyt\\AI智能投标\\能力组件-附件\\"
    image_paths = extract_all_docx_images(path, out_path)

    print("\n提取的图片列表:")
    for i, path in enumerate(image_paths, 1):
        print(f"{i}. {os.path.basename(path)}")



if __name__ == "__main__":
    main()